package cn.itcast.flink.operator;

import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.scala.typeutils.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import java.util.concurrent.TimeUnit;

/**
 * Author itcast
 * Date 2022/1/16 10:16
 * Desc 模拟消费kafka，每消费一条会 offset += 1，将 offset 保存到 checkpoint
 * 每产生 5 条数据，模拟 bug ，重启程序，接着上次最新的 offset 执行
 * 使用 OperatorState，数据结构是 ListState
 * 实现 OperatorState 需要实现 CheckpointedFunction 接口
 */
public class OperatorStateDemo {
    public static void main(String[] args) throws Exception {
        //1.创建流执行环境，设置并行度
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        //2.启用 checkpoint,每 1s 做一次checkpoint
        // 将 Operator State 保存到内存中， checkpoint 中
        // checkpoint 默认存在 JobManager 的内存
        // 设置 setStateStorage(HDFS)
        env.enableCheckpointing(1000);
        //3.设置重启策略，自动重启，重启3次，中间休眠 3s
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3,3000));
        //4.添加自定义数据源，将模拟的 offset 返回
        DataStreamSource<Long> source = env.addSource(new MySource());
        //5.map操作，将 Long 类型转换成字符串类型并打印输出
        source.printToErr();

        //6.执行流环境
        env.execute();
    }

    //自定义source实现 SourceFunction<Integer> 和 CheckpointedFunction
    //如果使用 Operator State ，实现 CheckpointedFunction
    public static class MySource implements SourceFunction<Long>, CheckpointedFunction {

        //1.定义变量
        //1.1定义标记用于循环生成标记
        volatile boolean isRunning = true;
        //1.2定义 currentCounter 用于保存当前的计数值
        Long currentCounter = 0L;
        //1.3定义 listState
        ListState<Long> listState = null;
        //1.4定义ListState描述
        ListStateDescriptor<Long> listStateDesc = new ListStateDescriptor("offsetState", Types.LONG());

        //2.重写 snapshotState 方法，清除状态和将最新的累加值添加到状态中
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            //将上一次的 state 清空
            listState.clear();
            listState.add(currentCounter);
        }

        //3.重写 initializeState 方法，获取状态并遍历Iterable，将其赋值给 currentCounter
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            //3.1 将历史存储的状态遍历出来 赋值给currentCounter变量
            listState = context.getOperatorStateStore().getListState(listStateDesc);
            //将从 state 中读取出来最新的 offset 赋值给 currentCounter
            Iterable<Long> longs = listState.get();
            //遍历这个集合将最新的值赋值给 currentCounter
            for (Long offset : longs) {
                currentCounter = offset;
            }
        }

        //4.重写 run方法，每秒循环收集累加的counter，每5个生成一个异常
        @Override
        public void run(SourceContext<Long> ctx) throws Exception {
            //4.1 持续循环
                while(isRunning) {
                    //4.2 对 currentCounter 累加
                    currentCounter ++;
                    //4.3 收集当前累加的值
                    ctx.collect(currentCounter);
                    //4.4 休眠 1 秒钟
                    TimeUnit.SECONDS.sleep(1);
                    //4.5 如果是 5 的倍数就模拟输出异常
                    if(currentCounter % 5==0){
                        throw new Exception("1/0 错误");
                    }
            }
        }

        //5.重写 cancel方法
        @Override
        public void cancel() {
            isRunning = false;
        }
    }
}
